from math import sqrt, pi, exp

data_set = [[3.393533211, 2.331273381, 0],
            [3.110073483, 1.781539638, 0],
            [1.343808831, 3.368360954, 0],
            [3.582294042, 4.67917911, 0],
            [2.280362439, 2.866990263, 0],
            [7.423436942, 4.696522875, 1],
            [5.745051997, 3.533989803, 1],
            [9.172168622, 2.511101045, 1],
            [7.792783481, 3.424088941, 1],
            [7.939820817, 0.791637231, 1]]


def separate_data_set(data_set):
    separated_data_set = dict()
    for i in range(len(data_set)):
        data = data_set[i]
        label = data[-1]
        if label not in separated_data_set:
            separated_data_set[label] = list()
        separated_data_set[label].append(data)
    return separated_data_set


def mean(columns):
    return sum(columns) / len(columns)


def standard_deviation(columns):
    avg = mean(columns)
    variance = sum([(n - avg) **2 for n in columns]) / len(columns)
    return sqrt(variance)


def summarize_data_set(data_set):
    summaries = list()
    summaries = [(mean(column), standard_deviation(column), len(column)) for column in zip(*data_set)]
    del(summaries[-1])
    return summaries


def summarize_by_class(data_set):
    summaries = dict()
    separated_data_set = separate_data_set(data_set)
    for label, rows in separated_data_set.items():
        summaries[label] = summarize_data_set(rows)
    return summaries


def calculate_probability(x, mean, std):
    exponent = exp(-(x-mean)**2 / (2 * std ** 2))
    return (1 / (sqrt(2 * pi) * std)) * exponent


def calculate_class_probabilities(summaries, test_data):
    total_rows = sum([summaries[label][0][2] for label in summaries])
    probabilities = dict()
    for class_value, summaries_values in summaries.items():
        probabilities[class_value] = summaries[class_value][0][2] / total_rows
        for i in range(len(summaries_values)):
            mean, std, count = summaries_values[i]
            probabilities[class_value] *= calculate_probability(test_data[i], mean, std)
    return probabilities


def predict(summaries, test_data):
    probabilities = calculate_class_probabilities(summaries, test_data)
    best_label, best_prob = None, -1
    for class_value, probability in probabilities.items():
        if best_label is None or probability > best_prob:
            best_prob = probability
            best_label = class_value
    return best_label


def naive_bayes(train_set, test_set):
    summarize = summarize_by_class(train_set)
    predictions = list()
    for test_data in test_set:
        y_pred = predict(summarize, test_data)
        predictions.append(y_pred)
    return (predictions)


if __name__ == '__main__':
    test_data = [[3.393533211, 2.331273381], [3.110073483, 1.781539638]]
    predictions = naive_bayes(data_set, test_data)
    print(predictions)
